import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# routes.py
from fastapi import APIRouter
from pydantic import BaseModel
from agent.research_agent import tool_agent
from langchain_core.messages import HumanMessage
from agent.research_assistant.analysts_graph import analystsGraph
from agent.research_assistant.interview_graph import interview_graph
from agent.research_assistant.report_graph import report_graph
from agent.research_assistant.state import Analyst
import json
from typing import Optional

# 创建路由器
router = APIRouter()

# 定义API请求模型
class InputModel(BaseModel):
    user_input: str
    thread_id: int
    analyst: Optional[Analyst] = None  # 设置默认值为None

@router.post("/partial/assistant")
async def get_assistant_response(input_model: InputModel):
    user_input = input_model.user_input
    thread_id = input_model.thread_id
    # Input
    max_analysts = 3
    topic = user_input
    thread = {"configurable": {"thread_id": thread_id}}
    event = analystsGraph.invoke({"topic":topic,"max_analysts":max_analysts,}, thread, stream_mode="values")
    analysts = event.get('analysts', [])
    
    # 将analysts转换为可序列化的字典列表
    analysts_json = []
    for analyst in analysts:
        analysts_json.append({
            "name": analyst.name,
            "affiliation": analyst.affiliation,
            "role": analyst.role,
            "description": analyst.description
        })
    
    # 返回JSON格式的响应
    return {"analysts": analysts_json}


@router.post("/partial/human_analyst_feedback")
async def human_analyst_feedback(input_model: InputModel):
    user_input = input_model.user_input
    thread_id = input_model.thread_id
    # Input
    thread = {"configurable": {"thread_id": thread_id}}
    analystsGraph.update_state(thread, {"human_analyst_feedback":
                            user_input}, as_node="human_feedback")

    event = analystsGraph.invoke(None, thread, stream_mode="values")
    print(event)
    analysts = event.get('analysts', [])
    
    # 将analysts转换为可序列化的字典列表
    analysts_json = []
    for analyst in analysts:
        analysts_json.append({
            "name": analyst.name,
            "affiliation": analyst.affiliation,
            "role": analyst.role,
            "description": analyst.description
        })
    
    # 返回JSON格式的响应
    return {"analysts": analysts_json}

@router.post("/partial/write_sections")
async def write_sections(input_model: InputModel):
    user_input = input_model.user_input
    thread_id = input_model.thread_id
    analyst = input_model.analyst
    # Input
    thread = {"configurable": {"thread_id": thread_id}}
    messages = [HumanMessage(f"So you said you were writing an article on {user_input}?")]
    interview = interview_graph.invoke({"analyst": analyst, "messages": messages, "max_num_turns": 2}, thread)
    return json.dumps(interview['sections'][0])

@router.post("/report/assistant")
async def get_report_assistant_response(input_model: InputModel):
    user_input = input_model.user_input
    thread_id = input_model.thread_id
    # Input
    max_analysts = 3
    topic = user_input
    thread = {"configurable": {"thread_id": thread_id}}
    event = report_graph.invoke({"topic":topic,"max_analysts":max_analysts,}, thread, stream_mode="values")
    analysts = event.get('analysts', [])
    
    # 将analysts转换为可序列化的字典列表
    analysts_json = []
    for analyst in analysts:
        analysts_json.append({
            "name": analyst.name,
            "affiliation": analyst.affiliation,
            "role": analyst.role,
            "description": analyst.description
        })
    
    # 返回JSON格式的响应
    return {"analysts": analysts_json}


@router.post("/report/human_analyst_feedback")
async def human_report_analyst_feedback(input_model: InputModel):
    user_input = input_model.user_input
    thread_id = input_model.thread_id
    # Input
    thread = {"configurable": {"thread_id": thread_id}}
    report_graph.update_state(thread, {"human_analyst_feedback":
                            user_input}, as_node="human_feedback")
    event = report_graph.invoke(None, thread, stream_mode="values")
    analysts = event.get('analysts', [])
    print(event)
    
    # 将analysts转换为可序列化的字典列表
    analysts_json = []
    for analyst in analysts:
        analysts_json.append({
            "name": analyst.name,
            "affiliation": analyst.affiliation,
            "role": analyst.role,
            "description": analyst.description
        })

    # 返回JSON格式的响应
    return {"analysts": analysts_json}

@router.post("/report/content")
async def human_report_analyst_feedback(input_model: InputModel):
    thread_id = input_model.thread_id
    # Input
    thread = {"configurable": {"thread_id": thread_id}}
    
    printState(thread)

    report_graph.update_state(thread, {"human_analyst_feedback":
                            None}, as_node="human_feedback")
    event = report_graph.invoke(None, thread, stream_mode="values")
    
    report = event.get('final_report', "")
    # 返回JSON格式的响应
    return {"report": report}


def printState(thread):
    state = report_graph.get_state(thread)
    print(state.values)